import csv
import os
import datetime
from src.core.dbhelper import DBHelper
from src.core.utils import *

class CsvImporter:
    """CSV导入器，负责从CSV文件导入各类数据"""
    
    def __init__(self, database_connector):
        """初始化CSV导入器"""
        self.db = database_connector
        self.db_helper = DBHelper(database_connector)
        
    def _determine_column_index(self, headers, column_param, default_column_name, default_index):
        """确定CSV文件中列的索引
        
        Args:
            headers: CSV文件的表头列表
            column_param: 用户指定的列参数（索引或名称）
            default_column_name: 默认的列名称
            default_index: 默认的列索引
        
        Returns:
            int: 列的索引，如果无法确定则返回None
        """
        # 如果用户指定了列参数
        if column_param:
            try:
                # 尝试将参数解析为索引
                index = int(column_param)
                return index
            except ValueError:
                # 参数不是索引，尝试作为列名查找
                for i, header in enumerate(headers):
                    if header.strip().lower() == column_param.strip().lower():
                        return i
                # 未找到匹配的列名
                return None
        
        # 如果用户未指定列参数，尝试通过默认列名查找
        for i, header in enumerate(headers):
            if header.strip().lower() == default_column_name.lower():
                return i
        
        # 如果只有一列，直接使用该列
        if len(headers) == 1:
            return 0
        
        # 使用默认索引
        return default_index if default_index < len(headers) else None
    
    def import_ips_from_csv(self, file_path, ip_column=None, owner_column=None):
        """从CSV文件导入IP地址"""
        # 检查文件是否存在
        if not os.path.exists(file_path):
            print(f"错误：文件 '{file_path}' 不存在")
            return 0, 1
        
        try:
            # 打开CSV文件
            with open(file_path, 'r', encoding='utf-8') as f:
                reader = csv.reader(f)
                headers = next(reader)  # 获取表头
                
                # 确定IP列的位置
                ip_col_index = self._determine_column_index(headers, ip_column, 'ip', 1)  # 默认第2列
                if ip_col_index is None:
                    print("错误：无法确定IP列")
                    return 0, 1
                
                # 确定单位名称列的位置
                owner_col_index = None
                if owner_column:
                    owner_col_index = self._determine_column_index(headers, owner_column, 'owner', -1)
                    if owner_col_index is None:
                        print(f"警告：无法确定单位名称列 '{owner_column}'，将不使用单位信息")
                
                # 导入数据
                imported_count = 0
                skipped_count = 0
                for row_num, row in enumerate(reader, 2):  # 行号从2开始（跳过表头）
                    if not row:  # 跳过空行
                        continue
                    
                    # 确保行有足够的列
                    if ip_col_index >= len(row):
                        print(f"警告：第 {row_num} 行数据不足，跳过")
                        skipped_count += 1
                        continue
                    
                    ip = row[ip_col_index].strip()
                    
                    # 验证IP格式
                    if not is_valid_ip(ip):
                        print(f"警告：第 {row_num} 行的IP地址 '{ip}' 格式无效，跳过")
                        skipped_count += 1
                        continue
                    
                    # 检查IP是否已存在
                    query = "SELECT * FROM ips WHERE ip = %s"
                    cursor = self.db.execute_query(query, (ip,))
                    if cursor and cursor.fetchone():
                        print(f"警告：IP地址 '{ip}' 已存在，跳过")
                        skipped_count += 1
                        continue
                    
                    # 获取单位信息
                    owner_id = None
                    owner_name = "无"
                    if owner_col_index is not None and owner_col_index < len(row):
                        csv_owner_name = row[owner_col_index].strip()
                        if csv_owner_name:
                            # 查找单位是否存在
                            query = "SELECT * FROM owners WHERE owner_name = %s"
                            cursor = self.db.execute_query(query, (csv_owner_name,))
                            owner = cursor.fetchone() if cursor else None
                            
                            if owner:
                                owner_id = owner['owner_id']
                                owner_name = owner['owner_name']
                            else:
                                # 创建新单位
                                current_time = datetime.datetime.now()
                                query = "INSERT INTO owners (owner_name, create_time, update_time) VALUES (%s, %s, %s)"
                                cursor = self.db.execute_query(query, (csv_owner_name, current_time, current_time))
                                if cursor:
                                    owner_id = cursor.lastrowid
                                    owner_name = csv_owner_name
                                    print(f"成功创建新单位: {csv_owner_name}, ID: {owner_id}")
                    
                    # 添加IP地址
                    current_time = datetime.datetime.now()
                    query = "INSERT INTO ips (ip, owner_id, owner_name, create_time, update_time) VALUES (%s, %s, %s, %s, %s)"
                    cursor = self.db.execute_query(query, (ip, owner_id, owner_name, current_time, current_time))
                    if cursor:
                        imported_count += 1
                        print(f"成功导入IP地址: {ip}")
                    else:
                        print(f"错误：导入IP地址 '{ip}' 失败")
                        skipped_count += 1
                
                print(f"\n导入完成：成功导入 {imported_count} 个IP地址，跳过 {skipped_count} 个")
                return imported_count, skipped_count
        except csv.Error as e:
            print(f"CSV文件解析错误: {str(e)}")
            return 0, 1
        except Exception as e:
            print(f"导入IP地址失败: {str(e)}")
            return 0, 1
    
    def import_domains_from_csv(self, file_path, domain_column=None, owner_column=None, default_owner_id=None):
        """从CSV文件导入域名"""
        # 检查文件是否存在
        if not os.path.exists(file_path):
            print(f"错误：文件 '{file_path}' 不存在")
            return 0, 1
        
        try:
            # 打开CSV文件
            with open(file_path, 'r', encoding='utf-8') as f:
                reader = csv.reader(f)
                headers = next(reader)  # 获取表头
                
                # 确定域名列的位置
                domain_col_index = self._determine_column_index(headers, domain_column, 'domain', 1)  # 默认第2列
                if domain_col_index is None:
                    print("错误：无法确定域名列")
                    return 0, 1
                
                # 确定单位名称列的位置
                owner_col_index = None
                if owner_column:
                    owner_col_index = self._determine_column_index(headers, owner_column, 'owner', -1)
                    if owner_col_index is None:
                        print(f"警告：无法确定单位名称列 '{owner_column}'，将不使用单位信息")
                
                # 导入数据
                imported_count = 0
                skipped_count = 0
                for row_num, row in enumerate(reader, 2):  # 行号从2开始（跳过表头）
                    if not row:  # 跳过空行
                        continue
                    
                    # 确保行有足够的列
                    if domain_col_index >= len(row):
                        print(f"警告：第 {row_num} 行数据不足，跳过")
                        skipped_count += 1
                        continue
                    
                    domain_name = row[domain_col_index].strip()
                    # 验证域名格式
                    if not is_valid_domain(domain_name):
                        print(f"警告：第 {row_num} 行的域名 '{domain_name}' 格式无效，跳过")
                        skipped_count += 1
                        continue
                    
                    # 获取单位信息
                    owner_id = default_owner_id
                    if owner_col_index is not None and owner_col_index < len(row):
                        owner = self.db_helper.get_or_add_owner(row[owner_col_index].strip())
                        if owner:
                            owner_id = owner['owner_id']

                    # 添加域名
                    try:
                        domain = self.db_helper.get_or_add_domain(domain_name, owner_id)
                        if '_is_new' in domain and domain['_is_new']:
                            imported_count += 1
                            print(f"成功导入域名: {domain_name}")
                        else:
                            print(f"警告：第 {row_num} 行的域名 '{domain_name}' 已存在，跳过")
                            skipped_count += 1
                    except Exception as e:
                        print(f"错误：导入域名 '{domain_name}' 失败: {str(e)}")
                        skipped_count += 1
                
                print(f"\n导入完成：成功导入 {imported_count} 个域名，跳过 {skipped_count} 个")
                return imported_count, skipped_count
        except csv.Error as e:
            print(f"CSV文件解析错误: {str(e)}")
            return 0, 1
        except Exception as e:
            print(f"导入域名失败: {str(e)}")
            return 0, 1
    
    def import_assets_from_csv(self, file_path, ip_column=None, domain_column=None, name_column=None, port_column=None, protocol_column=None, fingerprint_column=None, title_column=None, owner_column=None):
        """从CSV文件导入资产"""
        # 检查文件是否存在
        if not os.path.exists(file_path):
            print(f"错误：文件 '{file_path}' 不存在")
            return 0, 1
        
        try:
            # 打开CSV文件
            with open(file_path, 'r', encoding='utf-8') as f:
                reader = csv.reader(f)
                headers = next(reader)  # 获取表头
                
                # 确定各列的位置
                ip_col_index = self._determine_column_index(headers, ip_column, 'ip', -1)  # 不设置默认值
                domain_col_index = self._determine_column_index(headers, domain_column, 'domain', -1)  # 不设置默认值
                name_col_index = self._determine_column_index(headers, name_column, 'name', -1)  # 不设置默认值
                port_col_index = self._determine_column_index(headers, port_column, 'port', -1)  # 不设置默认值
                protocol_col_index = self._determine_column_index(headers, protocol_column, 'protocol', -1)  # 不设置默认值
                fingerprint_col_index = self._determine_column_index(headers, fingerprint_column, 'fingerprint', -1)  # 不设置默认值
                title_col_index = self._determine_column_index(headers, title_column, 'title', -1)  # 不设置默认值
                owner_col_index = self._determine_column_index(headers, owner_column, 'owner', -1)  # 不设置默认值
                
                # 验证必要的参数是否可用
                if port_col_index is None or protocol_col_index is None:
                    # 尝试通过默认的列名查找
                    port_col_index = self._determine_column_index(headers, None, 'port', -1)
                    protocol_col_index = self._determine_column_index(headers, None, 'protocol', -1)
                    
                    if port_col_index is None or protocol_col_index is None:
                        print("错误：必须能够找到端口列和协议列")
                        return 0, 1
                
                # 导入数据
                imported_count = 0
                skipped_count = 0
                for row_num, row in enumerate(reader, 2):  # 行号从2开始（跳过表头）
                    if not row:  # 跳过空行
                        continue
                    
                    # 确保行有足够的列
                    if port_col_index >= len(row) or protocol_col_index >= len(row):
                        print(f"警告：第 {row_num} 行数据不足，跳过")
                        skipped_count += 1
                        continue
                    
                    # 获取基本信息
                    name = row[name_col_index].strip() if name_col_index is not None and name_col_index < len(row) else None
                    ip = row[ip_col_index].strip() if ip_col_index is not None and ip_col_index < len(row) else None
                    domain = row[domain_col_index].strip() if domain_col_index is not None and domain_col_index < len(row) else None
                    port = row[port_col_index].strip()
                    protocol = row[protocol_col_index].strip()
                    fingerprint = row[fingerprint_col_index].strip() if fingerprint_col_index is not None and fingerprint_col_index < len(row) else None
                    title = row[title_col_index].strip() if title_col_index is not None and title_col_index < len(row) else None
                    
                    # 验证参数
                    if not port or not protocol:
                        print(f"警告：第 {row_num} 行缺少端口或协议，跳过")
                        skipped_count += 1
                        continue
                    
                    if not ip and not domain:
                        print(f"警告：第 {row_num} 行缺少IP和域名，跳过")
                        skipped_count += 1
                        continue
                    
                    # 验证IP格式
                    if ip and not is_valid_ip(ip):
                        print(f"警告：第 {row_num} 行的IP地址 '{ip}' 格式无效，跳过")
                        skipped_count += 1
                        continue
                    
                    # 验证域名格式
                    if domain and not is_valid_domain(domain):
                        print(f"警告：第 {row_num} 行的域名 '{domain}' 格式无效，跳过")
                        skipped_count += 1
                        continue
                    
                    # 检查资产是否已存在（通过IP、域名、端口、协议的组合）
                    query = "SELECT * FROM assets WHERE port = %s AND asset_protocol = %s"
                    params = [port, protocol]
                    
                    if ip:
                        query += " AND ip = %s"
                        params.append(ip)
                    if domain:
                        query += " AND domain = %s"
                        params.append(domain)
                    
                    cursor = self.db.execute_query(query, params)
                    if cursor and cursor.fetchone():
                        print(f"警告：资产已存在（IP={ip}, 域名={domain}, 端口={port}, 协议={protocol}），跳过")
                        skipped_count += 1
                        continue
                    
                    # 获取单位信息
                    owner_id = None
                    owner_name = "无"
                    if owner_col_index is not None and owner_col_index < len(row):
                        csv_owner_name = row[owner_col_index].strip()
                        if csv_owner_name:
                            # 查找单位是否存在
                            query = "SELECT * FROM owners WHERE owner_name = %s"
                            cursor = self.db.execute_query(query, (csv_owner_name,))
                            owner = cursor.fetchone() if cursor else None
                            
                            if owner:
                                owner_id = owner['owner_id']
                                owner_name = owner['owner_name']
                            else:
                                # 创建新单位
                                current_time = datetime.datetime.now()
                                query = "INSERT INTO owners (owner_name, create_time, update_time) VALUES (%s, %s, %s)"
                                cursor = self.db.execute_query(query, (csv_owner_name, current_time, current_time))
                                if cursor:
                                    owner_id = cursor.lastrowid
                                    owner_name = csv_owner_name
                                    print(f"成功创建新单位: {csv_owner_name}, ID: {owner_id}")
                    
                    # 添加资产
                    current_time = datetime.datetime.now()
                    query = "INSERT INTO assets (asset_name, ip, port, domain, asset_protocol, asset_tags, owner_id, create_time, update_time) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)"
                    
                    # 构建标签字段
                    tags = []
                    if fingerprint:
                        tags.append(f"fingerprint:{fingerprint}")
                    if title:
                        tags.append(f"title:{title}")
                    tags_str = ",".join(tags) if tags else None
                    
                    cursor = self.db.execute_query(query, (name, ip, port, domain, protocol, tags_str, owner_id, current_time, current_time))
                    if cursor:
                        imported_count += 1
                        print(f"成功导入资产: ID={cursor.lastrowid}")
                        if name:
                            print(f"  名称: {name}")
                        if ip:
                            print(f"  IP: {ip}")
                        if domain:
                            print(f"  域名: {domain}")
                        print(f"  端口: {port}")
                        print(f"  协议: {protocol}")
                        if tags_str:
                            print(f"  标签: {tags_str}")
                        if owner_id:
                            print(f"  所属单位: {owner_name} (ID: {owner_id})")
                    else:
                        print(f"错误：导入资产失败 (IP={ip}, 域名={domain}, 端口={port}, 协议={protocol})")
                        skipped_count += 1
                
                print(f"\n导入完成：成功导入 {imported_count} 个资产，跳过 {skipped_count} 个")
                return imported_count, skipped_count
        except csv.Error as e:
            print(f"CSV文件解析错误: {str(e)}")
            return 0, 1
        except Exception as e:
            print(f"导入资产失败: {str(e)}")
            return 0, 1
    
    def import_ipdomains_from_csv(self, file_path, ip_column=None, domain_column=None):
        """从CSV文件导入域名IP映射"""
        # 检查文件是否存在
        if not os.path.exists(file_path):
            print(f"错误：文件 '{file_path}' 不存在")
            return 0, 1
        
        try:
            # 打开CSV文件
            with open(file_path, 'r', encoding='utf-8') as f:
                reader = csv.reader(f)
                headers = next(reader)  # 获取表头
                
                # 确定IP列和域名列的位置
                ip_col_index = self._determine_column_index(headers, ip_column, 'ip', 0)  # 默认第1列
                domain_col_index = self._determine_column_index(headers, domain_column, 'domain', 1)  # 默认第2列
                
                if ip_col_index is None or domain_col_index is None:
                    print("错误：无法确定IP列或域名列")
                    return 0, 1
                
                # 导入数据
                imported_count = 0
                skipped_count = 0
                for row_num, row in enumerate(reader, 2):  # 行号从2开始（跳过表头）
                    if not row:  # 跳过空行
                        continue
                    
                    # 确保行有足够的列
                    if ip_col_index >= len(row) or domain_col_index >= len(row):
                        print(f"警告：第 {row_num} 行数据不足，跳过")
                        skipped_count += 1
                        continue
                    
                    ip = row[ip_col_index].strip()
                    domain = row[domain_col_index].strip()
                    
                    try:
                        ipdomain = self.db_helper.get_or_add_ipdomain(ip, domain)
                        if '_is_new' in ipdomain and ipdomain['_is_new']:
                            imported_count += 1
                            print(f"成功导入映射: 域名 {domain} <--> IP地址 {ip}")
                        else:
                            print(f"警告：域名 {domain} <--> IP地址 {ip} 已存在，跳过")
                            skipped_count += 1
                    except Exception as e:
                        print(f"错误：导入映射 '{domain} <--> {ip}' 失败。错误信息:{str(e)}")
                        skipped_count += 1
                
                print(f"\n导入完成：成功导入 {imported_count} 个域名IP映射，跳过 {skipped_count} 个")
                return imported_count, skipped_count
        except csv.Error as e:
            print(f"CSV文件解析错误: {str(e)}")
            return 0, 1
        except Exception as e:
            print(f"导入域名IP映射失败: {str(e)}")
            return 0, 1

    def import_owners_from_csv(self, file_path, name_column=None):
        """从CSV文件导入单位数据"""
        success_count = 0
        failed_count = 0
        
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                # 使用csv.reader读取CSV文件
                reader = csv.reader(f)
                # 获取表头
                headers = next(reader, None)
                if not headers:
                    print("错误: CSV文件为空")
                    return success_count, failed_count
                
                # 确定名称列的索引
                name_col_index = self._determine_column_index(headers, name_column, 'owner_name', 0)
                if name_col_index is None:
                    print("错误: 未找到单位名称列")
                    return success_count, failed_count
                
                # 处理每一行数据
                for row in reader:
                    # 跳过空行
                    if not row or all(cell.strip() == '' for cell in row):
                        continue
                    
                    # 确保行长度足够
                    if len(row) <= name_col_index:
                        print(f"错误: 行数据不足，跳过该行: {row}")
                        failed_count += 1
                        continue
                    
                    # 获取单位名称
                    owner_name = row[name_col_index].strip()
                    if not owner_name:
                        print("错误: 单位名称为空，跳过该行")
                        failed_count += 1
                        continue
                    
                    # 检查单位是否已存在
                    try:
                        sql = "SELECT owner_id FROM owners WHERE owner_name = %s"
                        cursor = self.db.execute_query(sql, (owner_name,))
                        
                        has_result = False
                        if hasattr(cursor, 'fetchone'):
                            has_result = cursor.fetchone() is not None
                        elif cursor:
                            has_result = len(cursor) > 0
                            
                        if has_result:
                            print(f"警告: 单位 '{owner_name}' 已存在，跳过")
                            failed_count += 1
                            continue
                        
                        # 添加单位
                        current_time = datetime.datetime.now()
                        sql = "INSERT INTO owners (owner_name, create_time, update_time) VALUES (%s, %s, %s)"
                        
                        cursor = self.db.execute_query(sql, (owner_name, current_time, current_time))
                        success_count += 1
                    except Exception as e:
                        print(f"导入单位 '{owner_name}' 失败: {str(e)}")
                        failed_count += 1
        except Exception as e:
            print(f"读取CSV文件时发生错误: {str(e)}")
            failed_count += 1
            
        return success_count, failed_count